package org.apache.flink.training.exercises.ridesandfares;

import org.apache.flink.api.common.JobExecutionResult;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.JoinFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.RichCoFlatMapFunction;
import org.apache.flink.streaming.api.functions.sink.PrintSinkFunction;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.training.exercises.common.datatypes.RideAndFare;
import org.apache.flink.training.exercises.common.datatypes.TaxiFare;
import org.apache.flink.training.exercises.common.datatypes.TaxiRide;
import org.apache.flink.training.exercises.common.source.TaxiFareGenerator;
import org.apache.flink.training.exercises.common.source.TaxiRideSourceGenerator;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * @Author hujf
 * @Date 2022/8/24 14:26
 * @Description TODO
 */
public class RidesAndFaresExercise {

    private final SourceFunction<TaxiRide> rideSource;
    private final SourceFunction<TaxiFare> fareSource;
    private final SinkFunction<RideAndFare> sink;

    public RidesAndFaresExercise(SourceFunction<TaxiRide> rideSource, SourceFunction<TaxiFare> fareSource, SinkFunction<RideAndFare> sink) {
        this.rideSource = rideSource;
        this.fareSource = fareSource;
        this.sink = sink;
    }

    public static void main(String[] args) throws Exception {

        RidesAndFaresExercise job =
                new RidesAndFaresExercise(
                        new TaxiRideSourceGenerator(),
                        new TaxiFareGenerator(),
                        new PrintSinkFunction<>());

        job.execute();
    }


    public JobExecutionResult execute() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();

        WatermarkStrategy<TaxiFare> strategy = WatermarkStrategy
                .<TaxiFare>forBoundedOutOfOrderness(Duration.ofHours(1))
                .withTimestampAssigner((fare, timestamp) -> fare.getEventTimeMillis());

        WatermarkStrategy<TaxiRide> strategy2 = WatermarkStrategy
                .<TaxiRide>forBoundedOutOfOrderness(Duration.ofHours(1))
                .withTimestampAssigner((fare, timestamp) -> fare.getEventTimeMillis());

        DataStream<TaxiRide> rides = environment.addSource(rideSource).filter(o -> o.isStart)
                .keyBy(o -> o.rideId).assignTimestampsAndWatermarks(strategy2);
        DataStream<TaxiFare> fare = environment.addSource(fareSource).keyBy(o -> o.rideId).assignTimestampsAndWatermarks(strategy);


        rides.join(fare)
                .where(data -> data.rideId)
                .equalTo(data -> data.rideId)
                .window(TumblingEventTimeWindows.of(Time.milliseconds(2)))
                .apply(new JoinFunction<TaxiRide, TaxiFare, RideAndFare>() {
                    @Override
                    public RideAndFare join(TaxiRide first, TaxiFare second) throws Exception {
                        return new RideAndFare(first, second);
                    }
                }).addSink(sink);

        //这里注意两个流只有键一致的时候才能连接。 keyBy 的作用是将流数据分区，当 keyed stream 被连接时，他们必须按相同的方式分区。这样保证了两个流中所有键相同的事件发到同一个实例上。这样也使按键关联两个流成为可能。
        //   rides.connect(fare).flatMap(new EnrichmentFunction()).addSink(sink);

        return environment.execute("Join Rides with Fares");
    }

    /**
     * 这个方法自动识别不同的键
     */
    private class EnrichmentFunction extends RichCoFlatMapFunction<TaxiRide, TaxiFare, RideAndFare> {

        private ValueState<TaxiRide> rideState;
        private ValueState<TaxiFare> fareState;

        @Override
        public void open(Configuration parameters) throws Exception {
            rideState = getRuntimeContext().getState(new ValueStateDescriptor<>("saved ride", TaxiRide.class));
            fareState = getRuntimeContext().getState(new ValueStateDescriptor<>("saved fare", TaxiFare.class));
        }

        //因为两个流的出入时间 不同 所以 可能进入1 时 2 已经有了   进入2时 1已经有了
        @Override
        public void flatMap1(TaxiRide ride, Collector<RideAndFare> out) throws Exception {

            TaxiFare fare = fareState.value();

            if (fare != null) {
                fareState.clear();
                out.collect(new RideAndFare(ride, fare));
            } else {
                rideState.update(ride);
            }
        }

        @Override
        public void flatMap2(TaxiFare fare, Collector<RideAndFare> out) throws Exception {
            TaxiRide ride = rideState.value();
            if (ride != null) {
                rideState.clear();
                out.collect(new RideAndFare(ride, fare));
            } else {
                fareState.update(fare);
            }
        }
    }
}
